FIX Transformers weight conversion PEFT model prefix - #3230

"""
Reproducer for scope_prefix mis-attribution in
`transformers.conversion_mapping.get_model_conversion_mapping`.

The function iterates `model.named_modules()` and sets
`transform.scope_prefix = module_name` on every non-root `PreTrainedModel` match.
This is correct when `model` *is* the topmost `PreTrainedModel`. But when `model`
is any wrapper that holds a `PreTrainedModel` at a non-root attribute path, the
wrapper path bleeds into `scope_prefix`. The resulting transforms then silently
refuse to match any state-dict key that doesn't carry that wrapper path, because
`_scoped_match` requires the key to start with `f"{scope_prefix}."`.

Nothing in this file imports peft. The wrapper is a plain `nn.Module` that stands
in for any external integration that nests a `PreTrainedModel` under some
attribute (PEFT, accelerate hooks, RLHF reference-model holders, custom
distillation setups, ...).

Proposed fix (transformers side, around the `scope_prefix` mechanism): compute
`scope_prefix` relative to the topmost `PreTrainedModel` encountered during the
walk, not relative to `model` itself. When the caller hands in a wrapper, the
topmost `PreTrainedModel` sits at some non-root path; re-basing makes its
`scope_prefix` stay unset (which is what we want), while still letting any
genuine nested sub-`PreTrainedModel`s get correctly scoped relative paths.
"""

import torch
from torch import nn

from transformers import ViTConfig, ViTModel
from transformers.conversion_mapping import get_model_conversion_mapping
from transformers.core_model_loading import WeightRenaming


class Wrapper(nn.Module):
    """Minimal external wrapper that holds a PreTrainedModel at a non-root path.

    Any integration that nests a PreTrainedModel like this triggers the bug.
    """

    def __init__(self, inner: ViTModel):
        super().__init__()
        self.inner = inner


# Tiny ViT; we never run a forward, only inspect named_modules() / call the
# conversion-mapping helper. ViTModel is convenient because its registered
# conversion mapping is a clean chain of `WeightRenaming`s (no MoE fusion noise).
config = ViTConfig(
    hidden_size=32,
    num_hidden_layers=2,
    num_attention_heads=4,
    intermediate_size=64,
)
with torch.device("meta"):  # avoid real allocation; structure is all we need
    vit = ViTModel(config)
wrapped = Wrapper(vit)


def renamings(conversions):
    return [t for t in conversions if isinstance(t, WeightRenaming)]


# ---------------------------------------------------------------------------
# 1. Baseline: call the helper with the actual PreTrainedModel as argument.
#    scope_prefix stays None on every transform — they apply globally.
# ---------------------------------------------------------------------------
direct = get_model_conversion_mapping(vit)
print("=== get_model_conversion_mapping(vit)  (unwrapped) ===")
for t in renamings(direct):
    print(
        f"  {t.source_patterns} -> {t.target_patterns}  "
        f"scope_prefix={t.scope_prefix!r}"
    )

# ---------------------------------------------------------------------------
# 2. Buggy: call the helper with a wrapper. The walk reports the ViTModel at
#    path 'inner', so every transform gets scope_prefix='inner'.
# ---------------------------------------------------------------------------
via_wrapper = get_model_conversion_mapping(wrapped)
print("\n=== get_model_conversion_mapping(Wrapper(vit)) ===")
for t in renamings(via_wrapper):
    print(
        f"  {t.source_patterns} -> {t.target_patterns}  "
        f"scope_prefix={t.scope_prefix!r}"
    )

# ---------------------------------------------------------------------------
# 3. Concrete consequence: a real v4-style ViT checkpoint key goes through the
#    rename chain cleanly via the unwrapped model but is silently left
#    unchanged via the wrapper (because `_scoped_match` rejects it for not
#    starting with 'inner.').
# ---------------------------------------------------------------------------
example_key = "encoder.layer.0.attention.attention.query.weight"


def apply_renamings(key, conversions):
    """Mimic the per-key loop transformers/core_model_loading uses."""
    for t in renamings(conversions):
        key, _ = t.rename_source_key(key)
    return key


renamed_direct = apply_renamings(example_key, direct)
renamed_wrapper = apply_renamings(example_key, via_wrapper)

print(f"\nkey: {example_key}")
print(f"  via unwrapped model -> {renamed_direct}")
print(f"  via wrapper          -> {renamed_wrapper}")

# The unwrapped path normalises encoder.layer. -> layers. and attention.query
# -> q_proj, as registered for ViTModel in conversion_mapping.py:120-128.
assert renamed_direct == "layers.0.attention.q_proj.weight", (
    f"sanity check failed; unwrapped rename produced {renamed_direct!r}"
)

# The wrapper path leaves the key untouched — that's the bug. scope_prefix
# 'inner' makes every renaming a no-op for keys that don't carry that prefix.
assert renamed_wrapper == example_key, (
    "expected the bug to manifest as a silent skip; "
    f"wrapper produced {renamed_wrapper!r}"
)

print(
    "\nBug reproduced: every transform returned via the wrapper has "
    "scope_prefix='inner', so the rename chain silently skips ordinary "
    "checkpoint keys.\n"
    "\nProposed fix sketch (in get_model_conversion_mapping):\n"
    "  1. Find the topmost PreTrainedModel once:\n"
    "       root_path = next(name for name, m in model.named_modules()\n"
    "                       if isinstance(m, PreTrainedModel))\n"
    "  2. Re-base each subsequent `module_name` against `root_path` before\n"
    "     deciding scope:\n"
    "       if module_name == root_path: module_name = ''\n"
    "       elif module_name.startswith(root_path + '.'):\n"
    "           module_name = module_name[len(root_path) + 1:]\n"
    "     The rest of the loop is unchanged. When `model` is itself a\n"
    "     PreTrainedModel, `root_path == ''` and the diff is a no-op."
)